from .io_utils import cond_mkdir
import os
from abc import abstractclassmethod

submit_options = [
    "image",
    "priority",
    "name",
    "num_gpu",
    "worker_cpu",
    "data_sources",
    "num_machine",
    "cluster",
    "monitor",
    "node_names",
    "num_cpu",
]


def parse_dlc_params(params_file):
    data = {}
    with open(params_file, "r") as f:
        for line in f.readlines():
            line = line.strip()
            if "command" in line:
                key = "command"
                val = line.split(f"{key}=")[-1]
            else:
                key, value = line.split("=")
                key = key.strip()
                val = value.strip()
            data.update({key: val})
    return data


def dump_dlc_params(data: dict, file: str):
    params = ""
    for key, value in data.items():
        if value:
            params += f"{key}={value}\n"
    with open(file, "w") as f:
        f.write(params)
    return file


def _read_single_exec(file, msg=""):
    with open(file, "r") as f:
        for i, line in enumerate(f.readlines()):
            # skip comment
            line = line.strip()
            if i == 0:
                if msg != "":
                    msg += " && "
            if line.startswith("#"):
                continue
            if line:
                msg += f"{line.strip()} && "
    msg += f"echo '[tjob] Success run {os.path.realpath(file)}' "
    return msg


def show_file(file):
    print("dlc params file: ", file)
    with open(file, "r") as f:
        for f in f.readlines():
            print(f.strip())


# [deprecated]
def read_exec_files(files: list):
    msg = ""
    for file in files:
        msg = _read_single_exec(file, msg)
    return msg


def dump_exec_file(cmd: str, file: str):
    cond_mkdir(file, True)
    with open(file, "w") as f:
        f.write("#!/bin/bash\n")
        f.write(cmd)


def _read_py_file(file):
    msg = ""
    namespace = {}
    exec(open(file).read(), namespace)
    namespace.pop("__builtins__")
    _has_run = False
    for k, v in namespace.items():
        if k == "run":
            _has_run = True
            if isinstance(v, str):
                msg += f"{v} && "
            elif isinstance(v, (list, tuple)):
                for cmd in v:
                    msg += f"{cmd} && "
            else:
                err = (
                    "The `type(run)` in pyfile only support [tuple, list, str]"
                    + f", but get {type(v)}"
                )
                raise ValueError(err)

    if _has_run is False:
        err = "`tjob` parse command in *.py file by identifying the `run` keyword.\n"
        err += "User should add `run = [your_run_command]` in *.py file, for example:\n"
        err += "    ...\n"
        err += "    run = ['ls', 'cd xxx', 'bash xxx.sh']\n"
        raise RuntimeError(err)

    msg += f"echo '[tjob] Success run {os.path.realpath(file)}' "
    return msg


def _read_py_file_dlc_kwargs(file):
    # there is some mappings
    namespace = {}
    dlc_kwargs = {}

    exec(open(file).read(), namespace)
    namespace.pop("__builtins__")
    for k, v in namespace.items():
        if k in submit_options:
            dlc_kwargs[k] = v
    return dlc_kwargs


def _read_yaml_file_dlc_kwargs(file):
    from .io_utils import load_yaml

    dlc_kwargs = load_yaml(file).get("submit", {})
    for k, v in dlc_kwargs.items():
        if k in submit_options:
            dlc_kwargs[k] = v
    return dlc_kwargs


class _SingleParser:
    # single file parser
    def __init__(self, file: str) -> None:
        self._file = file
        self._command = self.read()

    @property
    def command(self):
        return self._command

    @abstractclassmethod
    def read(self) -> str:
        # read file and concat to _command
        pass


class BshParser(_SingleParser):
    def read(self):
        return _read_single_exec(self._file)


class PyParser(_SingleParser):
    # TODO(yuezhang): repeat exec file, modify as once
    def read(self):
        return _read_py_file(self._file)

    @property
    def dlc_kwargs(self) -> dict:
        return _read_py_file_dlc_kwargs(self._file)


class YamlParser(_SingleParser):
    def read(self):
        from tjob.initialize.setup_envs import NEXUS

        # NOTE: here use thelp(a runner for yaml file).
        # Does it exists a better way to handle yaml?
        msg = (
            "echo 'exec yaml via thelp'"
            + f" && pip install thelp -i {NEXUS} "
            + f" && thelp run {os.path.realpath(self._file)} -vvv"
        )
        return msg

    @property
    def dlc_kwargs(self) -> dict:
        return _read_yaml_file_dlc_kwargs(self._file)


class CommonParser:
    def __init__(self, files) -> None:
        from tjob.initialize.setup_envs import deekeeper_enabled, DEEKEEPER_SCRIPT_PATH

        if files is None:
            self._command = "echo 'SEND NONE'"
            self._dlc_kwargs = {}
            return
        elif isinstance(files, str):
            files = [files]

        if deekeeper_enabled():
            files.insert(0, DEEKEEPER_SCRIPT_PATH)

        self._command = ""
        self._raw_cmd = ""
        self._dlc_kwargs = {}
        for i, file in enumerate(files):
            if file.endswith(".py"):
                # pyfile support args parse
                parser = PyParser(file)
                self._command += parser.command
                self._dlc_kwargs.update(parser.dlc_kwargs)

            elif file.endswith(".sh") or file.endswith(".bash"):
                self._command += BshParser(file).command

            elif file.endswith(".yaml") or file.endswith(".yml"):
                raise ValueError(
                    "As  *.py is MORE powerful than *.yaml when submit job.\n"
                    + "TJOB DONT support *.yaml submit anymore."
                    + "Plz replace as *.py file."
                )
                parser = YamlParser(file)
                self._command += parser.command
                self._dlc_kwargs.update(parser.dlc_kwargs)
            else:
                err = f"{file} must end with [.py, .sh, .bash, .yaml, .yml]"
                raise TypeError(err)

            # dlc cannot use \n
            if i != len(files) - 1:
                self._command += " && "

    @property
    def command(self) -> str:
        return self._command

    @property
    def dlc_kwargs(self) -> dict:
        return self._dlc_kwargs
